#!/usr/bin/env python
"""
Library related to FSA and RNA Design
"""

import os
import sys
import re
import string
from pyrna import utils
from pyrna import ocamlfsa
from pyrna.tlf import regexps
from pyrna.parser import parseRNAStructure

###############################################################################

ocamlFSA_PATH = ocamlfsa.OCAML_FSA_PATH

GrGFreqs_PATH = utils.which('grgfreqs') 
if GrGFreqs_PATH is None:
    sys.stderr.write("program grgfreqs needs to be in PATH.\n")
    sys.exit(1)
GrGFreqs_COMMAND = GrGFreqs_PATH + " --all %s %d -o %s 2>/dev/null"

###############################################################################

def has_base_constraint(seq):
    from pyrna.alphabet import unambiguous_rna_letters

    for b in seq:
        if b in unambiguous_rna_letters or b != 'N':
            return True

    return False

###############################################################################

def locateMotif(seq, motifs):
    """
    locate a set of motifs in a sequence
    """
    matches = {}
    for motif in motifs:
        iterator = re.compile(motif).finditer(seq)
        for mo in iterator:
            matches[mo.start()] = motif        

    return matches

###############################################################################

def findCommonMotifs(file1, file2):
    """Return common motifs between file1 and file2
    """
    ms1 = [line.rstrip() for line in open(file1).readlines()]
    ms2 = [line.rstrip() for line in open(file2).readlines()]

    comm = []
    for i in range(0, len(ms1)):
        for j in range(0, len(ms2)):
            if ms1[i] == ms2[j]:
                comm.append(ms1[i])
                break

    return comm

###############################################################################

def rnaFold(seq):
    """
    Fold a sequence by RNAfold
    """
    os.system(" echo " + seq + " | RNAfold > tmp")
    htmp = open("tmp", 'r')
    lines = htmp.readlines()
    htmp.close()
    os.system("rm tmp")
    
    res = lines[1].rstrip("\n")
    
    return (res.split()[0], re.compile('(\-?\d+\.\d+)').search(res).group(1))

###############################################################################

def chooseStrictPairedBase(base):
    """
    Choose a strict paired base. (AU GC)
    """
    if (base == 'A'):
        return 'U'
    if (base == 'U'):
        return 'A'
    if (base == 'G'):
        return 'C'
    if (base == 'C'):
        return 'G'

###############################################################################

def chooseWobblePairedBase(base):
    """
    Choose a wobble paired base. (AU GC GU)
    """
    if (base == 'A'):
        return 'U'
    if (base == 'U'):
        return 'AG'
    if (base == 'G'):
        return 'CU'
    if (base == 'C'):
        return 'G' 

###############################################################################

def isSeqAcceptedFsa(expSeq, constraintFsa):
    """
    Is seq accepted by a FSA
    """
    
    os.system(ocamlFSA_PATH + " accepts " + expSeq + " " + constraintFsa 
              + " > tmp")
    htmp = open("tmp", 'r')
    if htmp.readline().rstrip("\n") == "true":
        return True 
    htmp.close()
    os.system("rm tmp")
    
    return False
    
###############################################################################

def genFsaNotForbiddenMotifs(f_m, plot=False):
    """
    Generate FSA not containing any forbidden motifs
    """
    tick = utils.Tick()

    utils.debug('make prefix tree by calling make_pt2 ...')
    f_ptr = ocamlfsa.make_pt2(f_m)
    if plot: ocamlfsa.plotFsa(f_ptr, 'png')
    utils.info('ptr generated: %s' % f_ptr)
    utils.info('elapsed time: %s' % tick.tock())
    
    utils.debug('read prefix tree %s ...' % f_ptr)
    fsastruct = readFsa(f_ptr)
    utils.info('elapsed time: %s' % tick.tock())
        
    utils.debug('merge final states ...')
    fsastruct = mergeFinStateOfPTree(fsastruct)
    utils.info('elapsed time: %s' % tick.tock())
    
    utils.debug('add [AUCG]* to Start and Final states ...')
    f_fsa = f_m + '.fsa'
    fsastruct = addToStartFinLoop(fsastruct)
    writeFsa(fsastruct, f_fsa)
    if plot: ocamlfsa.plotFsa(f_fsa, 'png')
    utils.info('fsa for forbidden motifs generated: %s' % f_fsa)
    utils.info('elapsed time: %s' % tick.tock())
    
    f_fsa_min = ocamlfsa.minimize(f_fsa)
    if plot: ocamlfsa.plotFsa(f_fsa_min, 'png')
    utils.info('fsa minimized: %s' % f_fsa_min)
    utils.info('elapsed time: %s' % tick.tock())

    utils.info('complement and minimize ...')
    f_fsa_complement = ocamlfsa.complement(f_fsa_min)
    if plot: ocamlfsa.plotFsa(f_fsa_complement, "png")
    utils.info('fsa complemented: %s' % f_fsa_complement)
    utils.info('elapsed time: %s' % tick.tock())
    
    return f_fsa_complement

###############################################################################    

def designSeq(mdaMotif, expSeq, fbdnSetMotifs, outfile, plot=False):
    """
    design RNA sequences according mandatory motifs and forbidden motifs
    """
    # 1. FSA of expected sequence, language sequence (.Ls)
    f_Ls = outfile + ".Ls"
    """
    # method a:
    reLs  = seqToRe(expSeq)
    fsaLs =regexps.compile(reLs, minimize=True)    
    compileToFSA(str(fsaLs), f_Ls)
    """
    # method b: use direct transformation from seq to FSA
    seqToFsa(expSeq, f_Ls)

    # minimize 
    f_Ls_min = ocamlfsa.minimize(f_Ls)
    if plot: ocamlfsa.plotFsa(f_Ls_min, 'png')
    utils.info('Fsa for expected sequence generated: %s' % f_Ls_min)
    
    mdaMotifOutfile = outfile + '_' + mdaMotif
    # 2. L0: FSA contain mandatory motif only once 
    #    and no any other forbidden motifs
    # L0 = L1 inter L2
    # L1: FSA can contain M only once
    # L2: FSA can contain mandatory motif but not any other forbidden motifs
    f_L0 = mdaMotifOutfile + '.L0'
    if os.path.exists(f_L0):
        utils.info('L0 exists: %s' % f_L0)
    else:        
        f_L2 = genFsaForbiddenExceptMotif(fbdnSetMotifs, mdaMotif)
        if plot: ocamlfsa.plotFsa(f_L2, 'png')

        # start L1: FSA that contain M only once, (L1 = L3 . M . L4)
        # L3: words that don't end with a word of mdaLft, not containing M
    
        # f_L5: FSA containing M
        f_L5 = mdaMotifOutfile + ".L5"

        # f_L6: FSA not containing M
        f_L6 = mdaMotifOutfile + ".L6"

        if not os.path.exists(f_L5) or not os.path.exists(f_L6):
            fsa_L5 = regexps.compile("[AUCG]*" + mdaMotif + "[AUCG]*", 
                                     minimize=True)
            compileToFSA(str(fsa_L5), f_L5)
            ocamlfsa.complement(f_L5, f_L6)
        utils.info('L6 generated: %s' % f_L6)
        if plot: ocamlfsa.plotFsa(f_L6, 'png')
            
        # mdaLft: words that can complete M to its left to give M again
        # Not ending with those words mdaLft
        f_L3 = mdaMotifOutfile + ".L3"    
    
        f_mdaLft = mdaMotifOutfile + ".mdaLft"    
        mdaLft = genMdaLftWords(mdaMotif)
        utils.info('words that can complete M to its left to give M: %s'
                   % ','.join(mdaLft))
        
        reMdaLft = "[AUCG]*"
        if len(mdaLft) > 0:
            reMdaLft = reMdaLft + "(" + "|".join(mdaLft) + ")"
            fsaMdaLft = regexps.compile(reMdaLft, minimize=True)
            compileToFSA(str(fsaMdaLft), f_mdaLft)
            fsaMdaLft_cmpl = ocamlfsa.complement(f_mdaLft)
            if plot: ocamlfsa.plotFsa(fsaMdaLft_cmpl, 'png')
            utils.info('fsaMdaLft_cmpl generated: %s' % fsaMdaLft_cmpl)
            
            # inter with L6 to generate L3
            ocamlfsa.inter(fsaMdaLft_cmpl, f_L6, f_L3)    
        else:
            utils.info("No words can complete M to its left to give M")
            f_L3 = f_L6

        if plot: ocamlfsa.plotFsa(f_L3, 'png')            
        utils.info('L3 generated: %s' % f_L3)
        
        # L4: words that don't start with a word of mdaRgt
        # mdaRgt: words that can complete M to its right to give M again
        f_L4 = mdaMotifOutfile + ".L4"

        f_mdaRgt = mdaMotifOutfile + ".mdaRgt"
        mdaRgt = genMdaRgtWords(mdaMotif)
        utils.info('words that can complete M to its right to give M: %s'
                   % ','.join(mdaRgt))
        
        reMdaRgt = "[AUCG]*"
        if len(mdaLft) > 0:
            reMdaRgt = "(" + "|".join(mdaRgt) + ")" + reMdaRgt
            fsaMdaRgt = regexps.compile(reMdaRgt, minimize=True)
            compileToFSA(str(fsaMdaRgt), f_mdaRgt)
            fsaMdaRgt_cmpl = ocamlfsa.complement(f_mdaRgt)
            if plot: ocamlfsa.plotFsa(fsaMdaRgt_cmpl, 'png')
            utils.info('fsaMdaRgt_cmpl generated: %s' % fsaMdaRgt_cmpl)
                    
            # inter with L6 to generate L4
            ocamlfsa.inter(fsaMdaRgt_cmpl, f_L6, f_L4)   
        else:
            utils.info("No words can complete M to its right to give M")
            f_L4 = f_L6
            
        if plot: ocamlfsa.plotFsa(f_L4, 'png')            
        utils.info('L4 generated: %s' % f_L4)
        
        # L1 = L3 . M . L4
        f_L1 = mdaMotifOutfile + ".L1"
        concatFsa1MFsa2(f_L3, mdaMotif, f_L4, f_L1)
        utils.info('L1=L3+M+L4 generated: %s' % f_L1)

        f_L1_min = ocamlfsa.minimize(f_L1)
        if plot: ocamlfsa.plotFsa(f_L1_min, 'png')
        utils.info('L1 minimized: %s' % f_L1_min)
        
        # L0 = L1 inter L2
        ocamlfsa.inter(f_L1_min, f_L2, f_L0)
        if plot: ocamlfsa.plotFsa(f_L0, 'png') 
        utils.info('L0 generated: %s' % f_L0)
    # end L0

    # 3. L = Ls inter L0 
    ocamlfsa.inter(f_Ls_min, f_L0, outfile)
    if plot: ocamlfsa.plotFsa(outfile, 'png') 
    utils.info('Fsa generated: %s' % outfile)
        
###############################################################################

def genMdaLftWords(mdaMotif):
    """
    Calculate words that can complete M to its left to give M again
    """
    mdaLft = [] # words that can complete M to its left to give M again

    for i in range(1, len(mdaMotif)):
        #print i, mdaMotif[:i], mdaMotif[:len(mdaMotif)-i]  
        if mdaMotif[:i] + mdaMotif[:len(mdaMotif)-i] == mdaMotif:
            mdaLft.append(mdaMotif[:i])

    return mdaLft

###############################################################################

def genMdaLftFsa(mdaLft, outfile):
    """
    Construct FSA not ending with words 'mdaLft'
    """
    f = ["0!\n", "0,0,A\n", "0,0,C\n", "0,0,G\n", "0,0,U\n"]

    curr = 1 # start from 0
    for i in range(0, len(mdaLft)):
        # each word one path
        f.append("0," + str(curr) + "," + mdaLft[i][0] + "\n")
        for j in range(1, len(mdaLft[i])-1):
            f.append(str(curr) + "," + str(curr+1) + "," + mdaLft[i][j] + "\n")
            curr = curr + 1
        # to end for the last base    
        f.append(str(curr) + ",F," + mdaLft[i][-1] + "\n")
        curr = curr + 1
    f.append("!F\n")

    handle = open(outfile + ".1", 'w')
    for i in range(0, len(f)):
        handle.write(f[i])
    handle.close()
    
    os.system(ocamlFSA_PATH + " determinize -rn " + outfile + ".1" 
              + " > " + outfile + ".2")
    os.system(ocamlFSA_PATH + " complement -m -rn " + outfile + ".2" 
              + " > " + outfile)    

###############################################################################
    
def genMdaRgtWords(mdaMotif):
    """
    Calculate words that can complete M to its right to give M again
    """
    mdaRgt = [] # words that can complete M to its right to give M again

    for i in range(1, len(mdaMotif)):
        #print i, mdaMotif[i:], mdaMotif[0-i:]
        if mdaMotif[i:] + mdaMotif[0-i:] == mdaMotif:
            mdaRgt.append(mdaMotif[0-i:])

    return mdaRgt

###############################################################################

def genMdaRgtFsa(mdaRgt, outfile):
    """
    Construct FSA not starting with words 'mdaLft'
    """
    f = ["1!\n"]

    curr = 2 # start from 2 for internal states
    for i in range(0, len(mdaRgt)):
        # each word one path
        f.append("1," + str(curr) + "," + mdaRgt[i][0] + "\n")
        for j in range(1, len(mdaRgt[i])-1):
            f.append(str(curr) + "," + str(curr+1) + "," + mdaRgt[i][j] + "\n")
            curr = curr + 1
        # to end for the last base    
        f.append(str(curr) + ",0," + mdaRgt[i][-1] + "\n")
        curr = curr + 1
        
    f.append("0,0,A\n")
    f.append("0,0,C\n")
    f.append("0,0,G\n")
    f.append("0,0,U\n")
    f.append("!0\n") 

    handle = open(outfile + ".1", 'w')
    for i in range(0, len(f)):
        handle.write(f[i])
    handle.close()
    
    os.system(ocamlFSA_PATH + " determinize -rn " + outfile + ".1" 
              + " > " + outfile + ".2")
    os.system(ocamlFSA_PATH + " complement -m -rn " + outfile + ".2" 
              + " > " + outfile)    
    
###############################################################################

def genFsaForbiddenExceptMotif(forbiddenMotifSet, exceptMotif):
    """
    FSA of deleting mandatory motif from forbidden motif set
    """
    fbnm = readMotifSet(forbiddenMotifSet)
    fbnm_except = []
    for m in fbnm:
        if m not in exceptMotif:
            fbnm_except.append(m)
        else:
            utils.info('Drop %s from %s' % (m, forbiddenMotifSet))

    f_fbdnMotifs = forbiddenMotifSet + ".except." + exceptMotif
    writeMotifSet(fbnm_except, f_fbdnMotifs)
    utils.info('current forbidden motifs generated %s' % f_fbdnMotifs)
        
    # generate FSA
    f_fsaForbiddenExcept = genFsaNotForbiddenMotifs(f_fbdnMotifs)   
    
    utils.info('%s generated' % f_fsaForbiddenExcept)
    return f_fsaForbiddenExcept

###############################################################################

def isEmptyFsa(f_fsa):
    """
    Check whether a FSA is empty or not 
    """
    handle = open(f_fsa, 'r')
    initialLine = handle.readline()   
    handle.close()

    if initialLine == "" or initialLine[0:2] == "!!":
        print(f_fsa + " is empty")
        return True
    if initialLine == "<emptyset>!\n":
        print(f_fsa + " is empty")
        return True

    return False

###############################################################################

def concatFsa1MFsa2(fsa1, motif, fsa2, outfile):
    """
    Construct FSA for "fsa1 + M + fsa2"
    Currently not support len(motif) == 1 
    """   
    outfile_handle = open(outfile, 'w')
    
    f1mf2_starts = []
    f1mf2_prods  = []
    
    # read fsa1
    handle = open(fsa1, 'r')
    fsa1 = handle.readlines()
    handle.close()
    
    #if (fsa1[0][0:2] == "!!"):
    #    fsa1 = ["0!\n", "0,0,A\n", "0,0,C\n", "0,0,G\n", "0,0,U\n!", "0\n"]
    
    fsa1_starts = fsa1[0].rstrip("!\n").split(",")
    fsa1_ends   = fsa1[-1].lstrip("!").rstrip("\n").split(",")
    
    #print fsa1_starts, fsa1_ends
    # rename start states, and add them to the start states of concatenated FSA
    for i in range(0, len(fsa1_starts)):
        f1mf2_starts.append("tmp" + fsa1_starts[i])
    outfile_handle.write(",".join(f1mf2_starts) + "!")     

    # rename productions, and add them
    for i in range(1, len(fsa1)-1):
        prod = fsa1[i].rstrip("!\n").split(",")
        f1mf2_prods.append(["tmp"+prod[0], "tmp"+prod[1], prod[2]])                            
    
    # add prods: end state => "m1" via motif[0]  
    for i in range(0, len(fsa1_ends)):
        f1mf2_prods.append(["tmp"+fsa1_ends[i], "m1", motif[0]])

    # add prods for the motif, "m2" => "m3 via motif[2]
    for i in range(1, len(motif)-1):
        f1mf2_prods.append(["m"+str(i), "m"+str(i+1), motif[i]])
            
    # read fsa2
    handle = open(fsa2, 'r')
    fsa2 = handle.readlines()
    handle.close()
    
    #if (fsa2[0][0:2] == "!!"):
    #    fsa2 = ["0!\n", "0,0,A\n", "0,0,C\n", "0,0,G\n", "0,0,U\n!", "0\n"]

    # add "m(length(motif))" => fsa2 initial states via motif[-1]
    fsa2_starts = fsa2[0].rstrip("!\n").split(",")
    for i in range(0, len(fsa2_starts)):
        f1mf2_prods.append(["m"+str(len(motif)-1), fsa2_starts[i], motif[-1]])   
   
    for i in range(0, len(f1mf2_prods)):
        outfile_handle.write("\n" + ",".join(f1mf2_prods[i]))
    
    # write productions and final states
    outfile_handle.write("\n")
    for i in range(1, len(fsa2)):
        outfile_handle.write(fsa2[i])
    outfile_handle.close()
    
    #plotFsa(outfile, "pdf")
       
###############################################################################

def checkStruct(f_seqs, rd_parser):
    """
    Check structure by 'rd_parser' for sequences in file 'f_seqs'
    Using Stephane's parse module and nltk library
    """
    # Check structure for each generated sequence
    outfile = f_seqs + ".can"
    outhand = open(outfile, 'w')
    
    handle = open(f_seqs, 'r')
    while handle:
        currSeq = handle.readline().strip('\n')
        if currSeq == '':
            break
        
        print(currSeq)
        p = rd_parser.nbest_parse(seqToList(currSeq))
        if len(p) == 0: # could form the expStruct
            print("Not correct structure")
        else:
            print("Candidate")
            outhand.write(currSeq + "\n")
            
    handle.close()
    outhand.close()
    
###############################################################################

def compileToFSA(aStr, outfile):
    """
    Write a FSA from regexps.compile() to file in format accepted by FSA-OCAML
    """
    handle = open(outfile, 'w')
    # find initial states:
    stInitial = re.compile('initial state: (\d+)\n').findall(aStr)
    
    # find final states:
    stFinal = re.compile('state: (\d+) \(final\)').findall(aStr)
    #print stInitial, stInternal, stFinal

    handle.write(",".join(stInitial) + "!")

    currStart = -1
    for line in aStr.split("\n"):
        if re.compile('state: (\d+)').search(line):
            currStart = re.compile('state: (\d+)').search(line).group(1)
        elif re.compile('\s+(\w+) -> (\d+)').search(line):
            currSymbol = re.compile('\s+(\w+) -> (\d+)').search(line).group(1)
            currEnd    = re.compile('\s+(\w+) -> (\d+)').search(line).group(2)
            handle.write("\n" + currStart + "," + currEnd + "," + currSymbol)

    handle.write("!\n" + ",".join(stFinal) + "\n")

    handle.close()

###############################################################################

def seqToRe(seq):
    """
    Convert a sequence to a regular expression
    """
    print(seq)

    seqRe = []

    for i in range(0, len(seq)):
        if seq[i] == 'N':
            seqRe.append('[AUCG]')
        else:
            seqRe.append(seq[i])

    return "".join(seqRe)


###############################################################################

def seqToFsa(seq, outfile):
    """
    Convert a regular expression to a FSA
    Example seq: CCNUANNNNNNCC
    """
    handle = open(outfile, 'w')
    curr = 0
    handle.write(str(curr) + "!\n")
    for i in range(0, len(seq)):
        if seq[i] == 'N':
            handle.write(str(curr) + "," + str(curr+1) + ",A\n")
            handle.write(str(curr) + "," + str(curr+1) + ",U\n")
            handle.write(str(curr) + "," + str(curr+1) + ",C\n")
            handle.write(str(curr) + "," + str(curr+1) + ",G\n")
            curr = curr + 1
        else:
            handle.write(str(curr) + "," + str(curr+1) + "," + seq[i] + "\n")
            curr = curr + 1            
    # add final state
    handle.write("!" + str(curr) + "\n")
    
    handle.close()

###############################################################################

def motifSetToRe(mset):
    """
    Conver a set of motifs to a regular expression
    Usage: motifSetToRe(['CCC', 'GGG'])
    regexps.compile('A*(GGG|CCC)A', minimize=True).match('GGGA')
    """
    mre = []
    
    for i in range(len(mset)):
        #print i, mset[i]
        mre.append(mset[i])

    if (len(mset) == 0):
        return ""
    else:
        return "[AUCG]*(" + "|".join(mre) + ")[AUCG]*"

###############################################################################

def readMotifSet(infile):
    """
    Read a set of motifs from a file, one motif per line
    """

    ms = []
    
    handle = open(infile, 'r')
    while True:
        txt = handle.readline()
        if txt == '':
            break
        if txt[0] == '#':
            continue
        
        ms.append((txt.rstrip().split("\t"))[0])    
            
    handle.close()

    return ms

###############################################################################

def writeMotifSet(ms, outfile):
    """
    write a set of motifs to a file, one motif per line
    """
    
    handle = open(outfile, 'w')
    for i in range(0, len(ms)):
        handle.write(ms[i])
        handle.write("\n")

    handle.close()
    
###############################################################################

def grammarToCFG(grammar):
    """
    grammar from parseRNAStructure to CFG in nltk
    """
    cfg = []
    
    prods = grammar.split("\n")
    for i in range(1, len(prods)):
        prod = prods[i]
        prod = string.replace(prod, "(", "")
        prod = string.replace(prod, ")", "")
        prod = string.replace(prod, "'", "")
        prod = string.replace(prod, "    ", "")
        cfg.append(prod)
        
    cfg.reverse()
    cfg.append("TG -> 'G'")
    cfg.append("TC -> 'C'")
    cfg.append("TA -> 'A'")
    cfg.append("TU -> 'U'")
 
    return "\n".join(cfg)

###############################################################################

def seqToList(seq):
    """
    Convert a sequence to a list
    """
    return list(seq)

###############################################################################

def readFsa(fsaFile):
    """
    read a FSA
    """
    handle = open(fsaFile, 'r')
    
    lines = handle.readlines()
    iStates = lines[0].rstrip("!\n").split(",")
    fStates = lines[-1].lstrip("!").rstrip("\n").split(",")
    prods   = {}
    for i in range(1, (len(lines)-1)):
        prod = lines[i].rstrip("!\n").split(",")
    
        #print prod
        if prod[0] not in prods:
            prods[prod[0]] = {'prev':[], 'next':[]}
        if prod[1] not in prods:
            prods[prod[1]] = {'prev':[], 'next':[]}
            
        prods[prod[0]]['next'].append([prod[1], prod[2]])
        prods[prod[1]]['prev'].append([prod[0], prod[2]])

    handle.close()
    #print iStates, fStates, prods
    return [iStates, fStates, prods]

###############################################################################

def writeFsa(fsa, outfile):
    """
    write a FSA to a file
    """
    handle = open(outfile, 'w')
    handle.write(",".join(fsa[0]) + "!")
    for s, prods in fsa[2].items():
        #print s
        for i in range(0, len(prods['next'])):
            handle.write("\n" + s + "," + prods['next'][i][0] + "," 
                         + prods['next'][i][1])

    handle.write("!\n" + ",".join(fsa[1]) + "\n")

    handle.close()


###############################################################################

def mergeFinStateOfPTree(fsa):
    """
    Merge final states of a prefix tree
    """

    fStates = fsa[1]
    prods   = fsa[2]

    i = 0
    while i < len(fStates):
        if len(prods[fStates[i]]['next']) != 0:
            i = i + 1
            continue
        else: # exist internal terminal state
            # if 'FINAL' not exist
            if 'FINAL' not in prods:
                prods['FINAL'] = {'prev':[], 'next':[]}              
            
            for j in range(0, len(prods[fStates[i]]['prev'])):
                ps = prods[fStates[i]]['prev'][j][0]
                for k in range(0, len(prods[ps]['next'])):
                    if (prods[ps]['next'][k][0] == fStates[i]):
                        prods[ps]['next'][k][0] = 'FINAL'

            # del this final state
            del prods[fStates[i]]
            # del this final state from the set of final states
            del fStates[i]
    
    fStates.append('FINAL') # add to final states set
            
    return fsa

###############################################################################

def addToStartFinLoop(fsa):
    """
    Add [AUCG]* for start states
    Add [AUCG]* for final states [inner and terminal]
    """
    sStates = fsa[0]
    fStates = fsa[1]
    prods   = fsa[2]

    # Add the new start state with loop [AUCG]*
    prods['START'] = {'prev':[['START', 'A'],
                              ['START', 'U'],
                              ['START', 'C'],
                              ['START', 'G']], 
                      'next':[['START', 'A'],
                              ['START', 'U'],
                              ['START', 'C'],
                              ['START', 'G']]}

    # expect one start state
    assert len(sStates) == 1
    
    # process previous start states    
    while len(sStates) > 0:
        s = sStates.pop(0) # delete it
        sNext = prods[s]['next']
        
        for j in range(0, len(sNext)):
            prods['START']['next'].append(sNext[j]) # 
            
            # search in the 'prev' list for start 's', and replace it with 'START'
            for k in range(0, len(prods[sNext[j][0]]['prev'])):
                if prods[sNext[j][0]]['prev'][k][0] == s:
                    prods[sNext[j][0]]['prev'][k][0] = 'START'

        # delete the start state
        del prods[s]

    # add the new start state
    sStates.append('START')

    # Add loop for 'FINAL' state
    prods['FINAL']['next'].append(['FINAL', 'A'])
    prods['FINAL']['next'].append(['FINAL', 'U'])
    prods['FINAL']['next'].append(['FINAL', 'C'])
    prods['FINAL']['next'].append(['FINAL', 'G'])
    prods['FINAL']['prev'].append(['FINAL', 'A'])
    prods['FINAL']['prev'].append(['FINAL', 'U'])
    prods['FINAL']['prev'].append(['FINAL', 'C'])
    prods['FINAL']['prev'].append(['FINAL', 'G'])

    # For final states except 'FINAL', i.e., inner final states
    for i in range(0, len(fStates)):
        f = fStates[i]
        if (f == 'FINAL'):
            continue
        prods[f]['next'].append(['FINAL', 'A'])
        prods[f]['next'].append(['FINAL', 'U'])
        prods[f]['next'].append(['FINAL', 'C'])
        prods[f]['next'].append(['FINAL', 'G'])
        
        prods['FINAL']['prev'].append([f, 'A'])
        prods['FINAL']['prev'].append([f, 'U'])
        prods['FINAL']['prev'].append([f, 'C'])
        prods['FINAL']['prev'].append([f, 'G'])
        
    return fsa

###############################################################################

def compileMotifSetToFSA(forbiddenMS):
    """
    Don't use it for big motif set, because of time and memory usage.
    Construct FSA for a motif set by using regexps.compile regular expression
    [AUCG]*(AAAAAA|AAACTT|...)[AUCG]*
    
    One way to improve it:
      use Divide-and-Conquer to combine FSAs for smaller motif set.
      Here combine 100 (FSA_UNIT) to current FSA at one time, 
      which is better than combine one by one.  
    """
    
    FSA_UNIT = 100 
    fms = forbiddenMS[0:FSA_UNIT] # list
    reForbidden = motifSetToRe(fms)

    fsaForbidden = regexps.compile(reForbidden, minimize=True)
    compileToFSA(str(fsaForbidden), "forbidden.0.fsa")

    # FSA of forbidden motifs
    nextFsaUnion = ""
    for i in range(1, len(forbiddenMS)/FSA_UNIT + 1):

        currIndex = i * FSA_UNIT
        nextIndex = min((i+1)*FSA_UNIT, len(forbiddenMS))

        #print currIndex, nextIndex
        currFms = forbiddenMS[currIndex:nextIndex] # list
        currReForbidden = motifSetToRe(currFms)
        #print currReForbidden

        currFsaForbidden = regexps.compile(currReForbidden, minimize=True)
        currFileFsaForbidden = "forbidden." + str(i) + ".fsa"
        compileToFSA(str(currFsaForbidden), currFileFsaForbidden)
        
        if (i == 1):       
            currFsaUnion = "forbidden.0.fsa"
        else:    
            currFsaUnion = "union." + str(i) + ".fsa"
        nextFsaUnion = "union." + str(i+1) + ".fsa"
        os.system(ocamlFSA_PATH + " union " + currFsaUnion + " " 
                  + currFileFsaForbidden + " > tmp") # + fileFsaForbidden)
        os.system(ocamlFSA_PATH + " minimize tmp > tmp1")
        os.system(ocamlFSA_PATH + " rename tmp1 > " + nextFsaUnion)
        
    return nextFsaUnion

###############################################################################

def makePairTable(struct):
    """
    Make pair table for a secondary structure
    """
    pairedTable = []
    # Initialization
    for i in range(0, len(struct)):
        pairedTable.append(-1) # not paired: -1
    
    stack = []
    for i in range(0, len(struct)):
        if struct[i] == '.':
            pass
        elif struct[i] == '(':
            stack.append(i)
        elif struct[i] == ')':
            j = stack.pop()
            pairedTable[i] = j
            pairedTable[j] = i
        else:
            print("Wrong letter in structure\n")
    
    if not len(stack) == 0:
        print("Not correct structure!\n")
        
    return pairedTable

###############################################################################                 

def reverseComplementWobble(rna):
    """
    Generate a list of reverse complement of a rna sequences 
    with wobble base-pairings
    """
    bases = []
    for b in rna:
        bases.append(b)
    
    # reverse    
    bases.reverse() 
    
    rcs = [""]
    
    # complement
    for i in range(0, len(bases)):
        for j in range(0, len(rcs)):
            if (bases[i] == 'A'):
                rcs[j] = rcs[j] + 'U'
            if (bases[i] == 'U'):
                rcs.append(rcs[j] + 'G')
                rcs[j] = rcs[j] + 'A'
            if (bases[i] == 'G'):
                rcs.append(rcs[j] + 'U')
                rcs[j] = rcs[j] + 'C'
            if (bases[i] == 'C'):
                rcs[j] = rcs[j] + 'G'
                                                    
    return rcs    

###############################################################################                 

def validate_seq2struct(seq, struct):
    """
    Validate whether seq is compatible with struct.
    Return True if yes, otherwise False
    """
    if len(seq) != len(struct):
        return False
    
    pt = makePairTable(struct)
    for i in range(len(struct)):
        if pt[i] < 0: continue
        
        bi = seq[i]
        bj = seq[pt[i]]
        if bi not in chooseWobblePairedBase(bj):
            return False

    return True
    
###############################################################################

        
